# -*- coding: utf-8 -*-
import json
from openchemie import OpenChemIE
from PIL import Image
import os

# 这是一个辅助函数，用于处理结果中的图像对象，以便它们可以被序列化为 JSON
def convert_images_to_serializable(obj):
    if isinstance(obj, Image.Image):
        # 将 PIL Image 对象转换为一个占位符字符串或文件名
        # 这里我们简单地返回一个描述，实际应用中您可能希望保存图像并返回路径
        return f"PIL.Image object of size {obj.size}"
    if isinstance(obj, dict):
        return {k: convert_images_to_serializable(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [convert_images_to_serializable(i) for i in obj]
    return obj

def main():
    """
    这是一个演示如何使用 OpenChemIE 库的示例函数。
    它会初始化 OpenChemIE 提取器，处理一个示例 PDF 文件，
    然后将提取出的化学反应信息打印出来，并保存到一个 JSON 文件中。
    """
    print("Initializing OpenChemIE extractor...")
    # 初始化提取器。第一次运行时，它会自动从 Hugging Face Hub 下载所需的模型文件，
    # 这可能需要一些时间，具体取决于您的网络连接。
    extractor = OpenChemIE()
    print("Initialization complete.")

    # 定义输入的 PDF 文件路径和输出的 JSON 文件路径
    pdf_path = "example/acs.joc.2c00749.pdf"
    output_path = "extraction_results.json"

    # 检查示例 PDF 文件是否存在
    if not os.path.exists(pdf_path):
        print(f"Error: Example PDF not found at '{pdf_path}'")
        print("Please make sure you are running this script from the root directory of the OpenChemIE project.")
        return

    # 从 PDF 中提取化学反应信息
    print(f"\nExtracting chemical reactions from: {pdf_path}")
    # 调用核心提取方法
    results = extractor.extract_reactions_from_pdf(pdf_path)
    print("Extraction complete.")

    # 由于提取结果可能包含无法直接转换为 JSON 的 PIL Image 对象，
    # 我们需要先对结果进行处理。
    serializable_results = convert_images_to_serializable(results)

    # 将结果保存到 JSON 文件中，以便后续分析
    print(f"Saving results to {output_path}...")
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            # 使用 indent=2 来美化输出的 JSON 格式
            json.dump(serializable_results, f, ensure_ascii=False, indent=2)
        print(f"Successfully saved results to {output_path}")
    except Exception as e:
        print(f"Error saving results to JSON: {e}")

    # 打印提取结果的摘要信息
    print("\n--- Extraction Summary ---")
    if 'figures' in serializable_results and serializable_results['figures']:
        print(f"Found reactions in {len(serializable_results['figures'])} figures.")
    else:
        print("No reactions found in figures.")
    
    if 'text' in serializable_results and serializable_results['text']:
        print(f"Found reactions in text across {len(serializable_results['text'])} pages.")
    else:
        print("No reactions found in text.")
    print("------------------------")


if __name__ == "__main__":
    main()
